package com.example.controller;

import com.graduation.dto.AttackTaskAssignDO;
import com.graduation.mapper.AttackTaskAssignMapper;
import com.graduation.mapper.ExploitTaskMapper;
import com.graduation.mapper.LateralMovementMapper;
import com.graduation.mapper.VulnerabilityScanMapper;
import com.graduation.pojo.*;
import com.graduation.vo.AttackTaskAssignVO;
import com.graduation.vo.EmployeeVO;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.ResponseBody;

import javax.annotation.Resource;
import javax.servlet.http.HttpSession;
import java.util.stream.Collectors;

@Controller
@WebException
@RequestMapping("/attack-task-assign")
public class AttackTaskAssignController extends BaseController<AttackTaskAssignDO, AttackTaskAssignVO> {

    @Resource
    private AttackTaskAssignMapper attackTaskAssignMapper;
    @Resource
    private VulnerabilityScanMapper vulnerabilityScanMapper;
    @Resource
    private ExploitTaskMapper exploitTaskMapper;
    @Resource
    private LateralMovementMapper movementMapper;

    @ResponseBody
    @GetMapping("/get")
        public WebResult<AttackTaskAssignVO> get(Integer id) {
        return new WebResult<>(getVoByDo(attackTaskAssignMapper.selectById(id), AttackTaskAssignVO.class));
    }

    @ResponseBody
    @PostMapping("/page")
    public WebResult<PageResult<AttackTaskAssignVO>> page(AttackTaskAssignQuery query, HttpSession session) {
        EmployeeVO loginer = (EmployeeVO) session.getAttribute("loginer");
        if ("tester".equals(loginer.getRole())) {
            query.setExecutorEq(loginer.getName());
        }
        PageResult<AttackTaskAssignVO> result = new PageResult<>();
        result.setRowCount(attackTaskAssignMapper.selectCount(query));
        result.setPageCount((result.getRowCount() - 1) / query.getPageSize() + 1);
        if (query.getPageNum() > result.getPageCount()) {
            query.setPageNum(result.getPageCount());
        }
        result.setPageNum(query.getPageNum());
        // 设置分页查询偏移量
        query.setLimit((query.getPageNum() - 1) * query.getPageSize());
        result.setEntityList(attackTaskAssignMapper.selectList(query).stream()
                .map(e -> getVoByDo(e, AttackTaskAssignVO.class))
                .collect(Collectors.toList()));
        return new WebResult<>(result);
    }

    @ResponseBody
    @PostMapping("/reject")
    public WebResult<?> reject(Integer id) {
        AttackTaskAssignDO task = attackTaskAssignMapper.selectById(id);
//        if (!task.getCompleteStatus().equals("待接受")) {
//            return WebResult.failed("当前任务" + task.getCompleteStatus() + ",不能进行拒绝");
//        }
        task.setCompleteStatus("已拒绝");
        attackTaskAssignMapper.updateById(task);
        return WebResult.success();
    }

    @ResponseBody
    @PostMapping("/accept")
    public WebResult<?> accept(Integer id) {
        AttackTaskAssignDO task = attackTaskAssignMapper.selectById(id);
//        if (!task.getCompleteStatus().equals("待接受")) {
//            return WebResult.failed("当前任务" + task.getCompleteStatus() + ",不能进行接受");
//        }
        task.setCompleteStatus("已接受");
        attackTaskAssignMapper.updateById(task);
        return WebResult.success();
    }

    @ResponseBody
    @PostMapping("/complete")
    public WebResult<?> complete(Integer id) {
        AttackTaskAssignDO task = attackTaskAssignMapper.selectById(id);
//        if (!task.getCompleteStatus().equals("已接受")) {
//            return WebResult.failed("已接受的任务才能完成");
//        }
        task.setCompleteStatus("已完成");
        attackTaskAssignMapper.updateById(task);
        return WebResult.success();
    }

    @RequestMapping("/to-attack-task-assign")
    public String toAttackTaskAssign() {
        return "attack/attackTaskAssignAdmin";
    }

    @RequestMapping("/to-view-attack-task-assign")
    public String toViewAttackTaskAssign() {
        return "attack/attackTaskAssignView";
    }

    @RequestMapping("/to-operate")
    public String toOperate(Model model, Integer id) {
        model.addAttribute("id", id);
        return "attack/operateAttackTaskAssign";
    }

    @Override
    protected void buildVo(AttackTaskAssignDO entity, AttackTaskAssignVO vo) {
        VulnerabilityScanQuery scanQuery = new VulnerabilityScanQuery();
        scanQuery.setRelatedTask(entity.getId());
        vo.setScanNum(vulnerabilityScanMapper.selectCount(scanQuery));
        ExploitTaskQuery exploitTaskQuery = new ExploitTaskQuery();
        exploitTaskQuery.setRelatedTask(entity.getId());
        vo.setExploitNum(exploitTaskMapper.selectCount(exploitTaskQuery));
        LateralMovementQuery movementQuery = new LateralMovementQuery();
        movementQuery.setRelatedTask(entity.getId());
        vo.setMovementNum(movementMapper.selectCount(movementQuery));
    }
}
